import os
import os.path as osp

import six
from yacs.config import CfgNode as CN

_C = CN()

cfg = _C

_C.SYSTEM = CN()

_C.SYSTEM.GPUS_INDEX = '0'
_C.SYSTEM.NUM_WORKERS = 4

_C.TRAIN = CN()

_C.TRAIN.MODEL_NAME = ''
_C.TRAIN.N_CLASS = 10

_C.TRAIN.DATA_PATH = 'data'
_C.TRAIN.LOG_PATH = 'results/log'
_C.TRAIN.OUTPUT_DIR = 'result'
_C.TRAIN.SAVE_CKPT_PATH = 'results/ckpt'
_C.TRAIN.LOAD_CKPT_PATH = 'results/ckpt'

_C.TRAIN.MIN_EOPCH = 5
_C.TRAIN.ITER_EOPCH = 50

# basic hyperparamer 
_C.TRAIN.EPOCHS = 48
_C.TRAIN.BATCH_SIZE = 16
_C.TRAIN.LEARNING_RATE = 3e-4

# SOLVER paramer
_C.TRAIN.SOLVER = CN()
_C.TRAIN.SOLVER.MOMENTUM = 0.9
_C.TRAIN.SOLVER.GAMMA = 0.2
_C.TRAIN.SOLVER.WEIGHT_DECAY = 5e-4

# lr scheduler
_C.TRAIN.LR_SCHEDULER = CN()
_C.TRAIN.LR_SCHEDULER.STEP_SIZE = 5
_C.TRAIN.LR_SCHEDULER.T_STRAT = 3
_C.TRAIN.LR_SCHEDULER.T_MULT = 2
_C.TRAIN.LR_SCHEDULER.LAST_EPOCH = -1

# SWA option
_C.TRAIN.SWA = CN()
_C.TRAIN.SWA.SWA_START = 42
_C.TRAIN.SWA.SWA_FREQ = 5
_C.TRAIN.SWA.SWA_LR_RATE = 0.6 # times with lr
# _C.TRAIN.SCALES = (2, 4, 8, 16)

def get_output_dir(datasets, training=True):
    """Get the output directory determined by the current global config."""
    assert isinstance(datasets, tuple([tuple, list] + list(six.string_types))), \
        'datasets argument must be of type tuple, list or string'
    is_string = isinstance(datasets, six.string_types)
    dataset_name = datasets if is_string else ':'.join(datasets)
    tag = 'train' if training else 'test'
    # <output-dir>/<train|test>/<dataset-name>/<model-type>/
    outdir = osp.join(_C.TRAIN.OUTPUT_DIR, tag, dataset_name)
    if not osp.exists(outdir):
        os.makedirs(outdir)
    return outdir
